package com.pig4cloud.pigx.knowledge.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.pig4cloud.pigx.common.core.constant.enums.YesNoEnum;
import com.pig4cloud.pigx.knowledge.config.properties.AiKnowledgeProperties;
import com.pig4cloud.pigx.knowledge.dto.AiDocumentDTO;
import com.pig4cloud.pigx.knowledge.entity.AiDatasetEntity;
import com.pig4cloud.pigx.knowledge.entity.AiDocumentEntity;
import com.pig4cloud.pigx.knowledge.entity.AiSliceEntity;
import com.pig4cloud.pigx.knowledge.mapper.AiDatasetMapper;
import com.pig4cloud.pigx.knowledge.mapper.AiDocumentMapper;
import com.pig4cloud.pigx.knowledge.service.AiAssistantService;
import com.pig4cloud.pigx.knowledge.service.AiDocumentService;
import com.pig4cloud.pigx.knowledge.service.AiSliceService;
import com.pig4cloud.pigx.knowledge.support.constant.SourceTypeEnums;
import com.pig4cloud.pigx.knowledge.support.constant.SummaryStatusEnums;
import com.pig4cloud.pigx.knowledge.support.handler.source.FileSourceTypeHandler;
import com.pig4cloud.pigx.knowledge.support.handler.source.UploadSourceTypeHandler;
import com.pig4cloud.pigx.knowledge.support.provider.ModelProvider;
import com.pig4cloud.pigx.knowledge.support.util.PromptBuilder;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 知识文档
 *
 * @author pig
 * @date 2024-03-14 13:38:59
 */
@Service
@Slf4j
@RequiredArgsConstructor
public class AiDocumentServiceImpl extends ServiceImpl<AiDocumentMapper, AiDocumentEntity>
		implements AiDocumentService {

	private final List<FileSourceTypeHandler> sourceTypeHandlerList;

	private final UploadSourceTypeHandler uploadSourceTypeHandler;

	private final AiKnowledgeProperties knowledgeProperties;

	private final AiDatasetMapper datasetMapper;

	private final AiSliceService sliceService;

	private final ModelProvider modelProvider;

	/**
	 * 获取文档页面
	 * @param page 页
	 * @param aiDocument AI 文档
	 * @return {@link Page }<{@link AiDocumentEntity }>
	 */
	@Override
	public Page<AiDocumentEntity> getDocumentPage(Page<AiDocumentEntity> page, AiDocumentEntity aiDocument) {
		return baseMapper.getDocumentPage(page, aiDocument);
	}

	/**
	 * 保存文档
	 * @param handler
	 * @param aiDocumentDTO 文档传输对象
	 * @param file
	 * @return true/false
	 */
	@Override
	@Async
	public void save(FileSourceTypeHandler handler, AiDocumentDTO aiDocumentDTO) {
		handler.handle(aiDocumentDTO);
	}

	/**
	 * 通过id删除文档和切片
	 * @param idList id 列表
	 * @return R
	 */
	@Override
	@Transactional(rollbackFor = Exception.class)
	public Boolean removeDocumentAndSliceBatchByIds(List<Long> idList) {
		// 删除切片 and 向量
		List<Long> sliceIdList = sliceService
			.list(Wrappers.<AiSliceEntity>lambdaQuery().in(AiSliceEntity::getDocumentId, idList))
			.stream()
			.map(AiSliceEntity::getId)
			.collect(Collectors.toList());

		if (CollUtil.isNotEmpty(sliceIdList)) {
			sliceService.removeSliceAndEbeddingById(sliceIdList);
		}

		// 删除文档
		baseMapper.deleteBatchIds(idList);
		return Boolean.TRUE;
	}

	/**
	 * 重试文档
	 * @param documentDTO 文档 DTO
	 */
	@SneakyThrows
	@Override
	@Async
	public void retryDocument(AiDocumentDTO documentDTO) {
		AiDocumentEntity documentEntity = baseMapper.selectById(documentDTO.getId());
		// 重新处理文件
		documentDTO.setFiles(List.of(documentEntity));
		uploadSourceTypeHandler.handle(documentDTO);
	}

	/**
	 * 重试问题
	 * @param aiDocument AI 文档
	 */
	@Override
	@Async
	public void retryIssue(AiDocumentDTO documentDTO) {
		AiDocumentEntity documentEntity = baseMapper.selectById(documentDTO.getId());
		JSONObject documentConfigObj = JSONUtil.parseObj(documentEntity.getDocumentConfig());
		documentEntity.setFileType("md");
		documentDTO.setRepoType(documentConfigObj.getStr(AiDocumentDTO.Fields.repoType));
		documentDTO.setRepoOwner(documentConfigObj.getStr(AiDocumentDTO.Fields.repoOwner));
		documentDTO.setRepoName(documentConfigObj.getStr(AiDocumentDTO.Fields.repoName));
		documentDTO.setAccessToken(documentConfigObj.getStr(AiDocumentDTO.Fields.accessToken));
		documentDTO.setFiles(List.of(documentEntity));

		sourceTypeHandlerList.stream()
			.filter(handler -> handler.supports(documentDTO.getSourceType()))
			.findFirst()
			.ifPresent(handler -> handler.handle(documentDTO));
	}

	@Override
	public void summaryDocument(AiDocumentEntity documentEntity) {
		AiDatasetEntity aiDataset = datasetMapper.selectById(documentEntity.getDatasetId());
		AiAssistantService assistantService = modelProvider.getAiAssistant(aiDataset.getSummaryModel()).getValue();

		// 如果未开启文档总结，则跳过
		if (YesNoEnum.NO.getCode().equals(aiDataset.getPreSummary())) {
			return;
		}

		// 如果是 QA 文档则跳过总结
		if (SourceTypeEnums.QA.getType().equals(documentEntity.getSourceType())) {
			return;
		}

		// 查询文档下所有数据,如果文档下没有未切片的数据，则跳过
		List<AiSliceEntity> sliceEntityList = sliceService
			.list(Wrappers.<AiSliceEntity>lambdaQuery().eq(AiSliceEntity::getDocumentId, documentEntity.getId()));

		if (sliceEntityList.isEmpty()) {
			return;
		}

		// 当前文档下的所有未训练数据，拼接成一个文档
		String documentContent = sliceEntityList.stream()
			.map(AiSliceEntity::getContent)
			.collect(Collectors.joining(StrUtil.LF));

		try {
			String result = assistantService
				.chat(PromptBuilder.render("knowledge-rag-summary.st", Map.of(AiDocumentEntity.Fields.summary,
						StrUtil.subSufByLength(documentContent, knowledgeProperties.getMaxSummary()))));
			documentEntity.setSummary(result);
			documentEntity.setSummaryStatus(SummaryStatusEnums.SUMMARYED.getStatus());
		}
		catch (Exception e) {
			log.warn("文档 {} 总结失败", documentEntity.getName(), e);
			documentEntity.setSummaryFailReason(e.getMessage());
			documentEntity.setSummaryStatus(SummaryStatusEnums.FAILED.getStatus());
		}

		baseMapper.updateById(documentEntity);
	}

}
